import sounddevice as sd  # 导入 sounddevice 库，用于录音
import numpy as np  # 导入 NumPy 库，用于处理音频数据
from pathlib import Path  # 导入 Path 模块，用于处理文件路径
import time  # 导入 time 模块，用于延时操作
from datetime import datetime  # 导入 datetime 模块，用于生成时间戳
from scipy.io.wavfile import write  # 导入 write 函数，用于保存 WAV 文件


class AudioCollector:  # 定义 AudioCollector 类，用于收集音频样本
    def __init__(self, root_dir="dataset"):  # 初始化方法，接收根目录参数
        self.root_dir = Path(root_dir)  # 将根目录转换为 Path 对象
        self.sample_rate = 16000  # 设置采样率为 16kHz
        self.duration = 3  # 设置每个样本的录音时长为 3 秒

        # 创建目录结构
        (self.root_dir / "train/wake").mkdir(parents=True, exist_ok=True)  # 创建 wake 子目录（如果不存在）
        (self.root_dir / "train/not_wake").mkdir(parents=True, exist_ok=True)  # 创建 not_wake 子目录

    def _record(self, save_path):  # 定义私有方法 _record，用于执行单次录音
        """执行单次录音"""
        print("开始录音...（说话！）")  # 提示用户开始录音
        audio = sd.rec(int(self.duration * self.sample_rate),  # 录制指定时长的音频
                       samplerate=self.sample_rate,  # 设置采样率
                       channels=1,  # 单声道录音
                       blocking=True)  # 阻塞模式，等待录音完成
        sd.wait()  # 确保录音完成
        # 保存为 16bit PCM 格式
        write(save_path, self.sample_rate, (audio * 32767).astype(np.int16))  # 将音频保存为 WAV 文件
        print(f"保存到 {save_path}")  # 打印保存路径

    def collect_wake_word(self, count=100):  # 定义方法 collect_wake_word，用于采集唤醒词样本
        """采集唤醒词样本"""
        for i in range(count):  # 循环采集指定数量的样本
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")  # 生成时间戳
            save_path = self.root_dir / f"train/wake/{timestamp}_{i}.wav"  # 构造保存路径
            self._record(save_path)  # 调用 _record 方法进行录音
            time.sleep(1)  # 每次录音后等待 1 秒

    def collect_background(self, count=300):  # 定义方法 collect_background，用于采集背景音样本
        """采集背景音样本"""
        input("请保持环境安静，准备采集背景音...（按回车开始）")  # 提示用户准备采集背景音
        for i in range(count):  # 循环采集指定数量的样本
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")  # 生成时间戳
            save_path = self.root_dir / f"train/not_wake/{timestamp}_{i}.wav"  # 构造保存路径
            self._record(save_path)  # 调用 _record 方法进行录音
            time.sleep(0.3)  # 每次录音后等待 0.3 秒


if __name__ == "__main__":  # 主程序入口
    collector = AudioCollector()  # 创建 AudioCollector 实例

    print("1. 采集唤醒词（需要说唤醒词）")  # 提示用户采集唤醒词
    collector.collect_wake_word(count=50)  # 采集 50 个唤醒词样本

    print("\n2. 采集背景音（保持安静或日常环境声）")  # 提示用户采集背景音
    collector.collect_background(count=150)  # 采集 150 个背景音样本
    #这段代码实现了一个简单的音频采集工具，用于收集唤醒词和背景音样本。它通过 sounddevice 录制音频，并将结果保存为 WAV 文件，存储在指定目录中。